#!/usr/bin/env python
# coding=utf-8

from pymongo import MongoClient
import tushare as ts
import numpy as np
import pandas as pd
import json

# check一下是否因为stocks_list不全导致一些股票数据没有入库
# 如果没有入库的话再抓取下
def get_db():
    client = MongoClient("localhost", 27017)
    db = client.stock
    return db

if __name__ == "__main__":
    db = get_db()
    print(db)

    with open("stock_list.txt") as f:
        for line in f:
            line = line.strip()
            if line == '':
                continue
            print(line)
            t = db.stocks.find_one({'code': line})
            if t == None:
                print(line + "is missing, getting data")
                origin_data = ts.get_k_data(line,  ktype='D', start='2005-01-01')
                if origin_data.empty:
                    print(line + ": data is empty")
                    continue
                db.stocks.insert(json.loads(origin_data.to_json(orient='records', date_format='iso')))
                print(line + ": data insert ok!")